/*
Usage:
Build 2d map from a point cloud file and a path file, and save the map as an image.
**NOTE**: All paths which not start with `/` are relative to `ws_618/src/map_manager/`

- Run directly with default parameters:
ros2 launch map_manager build_2d_map.launch.py

- OR specify parameters: (use ror denoise cloud)
ros2 launch map_manager build_2d_map.launch.py \
    pcl_file_path:=maps/lab/scans.pcd \
    path_file_path:=maps/lab/path.csv \
    flag_build_2d_map:=true \
    voxel_leaf_size:=0.1 \
    ror_radius:=0.05 \
    ror_min_neighbors:=30 \
    min_pz_gain:=1.3 \
    max_pz_gain:=1.0 \
    grid_value_threshold:=50.0

(no denoise)
ros2 launch map_manager build_2d_map.launch.py \
    pcl_file_path:=maps/lab/scans.pcd \
    path_file_path:=maps/lab/path.csv \
    flag_build_2d_map:=true \
    voxel_leaf_size:=0.1 \
    ror_radius:=-1.0 \
    min_pz_gain:=0.0 \
    max_pz_gain:=0.5 \
    grid_value_threshold:=90.0

- Load a 2D map from a PNG file and publish it as an OccupancyGrid message.
ros2 launch map_manager build_2d_map.launch.py \
    pcl_file_path:=relative_path/to/your/file.pcd \
    path_file_path:=relative_path/to/your/file.csv \
    flag_build_2d_map:=false \
    map_file_path:=relative_path/to/your/map.png
*/

#include <filesystem>
#include "rclcpp/rclcpp.hpp"
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <iostream>
#include <sensor_msgs/msg/point_cloud2.h>
#include <pcl_conversions/pcl_conversions.h>
#include <nav_msgs/msg/occupancy_grid.hpp>
#include <map_manager/utils.hpp>

using point_cloud_xyz = pcl::PointCloud<pcl::PointXYZ>;

double grid_value_threshold = 50.0; // Threshold for occupancy grid, cells with value > threshold will be marked as occupied
double resolution = 0.1; // 2d map resolution: 10 cm per cell
double min_pz = 1e9, max_pz = -1e9, min_pz_gain = 1.0, max_pz_gain = 1.0;
double voxel_leaf_size = 0.1;  // Voxel grid filter leaf size for downsampling the point cloud
double ror_radius = 0.05;  // Radius for radius outlier removal
int ror_min_neighbors = 30; // Minimum number of neighbors for radius outlier removal
bool flag_build_2d_map, reverse;
std::string root_dir = std::string(ROOT_DIR);

void update_min_max_pz_from_path_file(const std::string &path_file_path) {
    // Implement CSV file loading logic here if needed
    std::string line;
    std::ifstream file(path_file_path);
    if (!file.is_open()) {
        RCLCPP_ERROR(rclcpp::get_logger("Build2DMapNode"), "Failed to open path file: %s", path_file_path.c_str());
        return;
    } else {
        RCLCPP_INFO(rclcpp::get_logger("Build2DMapNode"), "Loading path file from: %s", path_file_path.c_str());
    }

    // Read the first line to find the index of the 'pz' column
    std::getline(file, line);
    int pz_col_index = -1;
    std::istringstream ss(line);
    std::string item;
    while (std::getline(ss, item, ',')) {
        ++pz_col_index;
        if (item == "pz") break;
    }
    if (pz_col_index == -1) {
        RCLCPP_ERROR(rclcpp::get_logger("Build2DMapNode"), "Column 'pz' not found in path file: %s", path_file_path.c_str());
        return;
    }

    // Read the rest of the file to find min and max pz values
    while (std::getline(file, line)) {
        std::istringstream ss(line);
        double pz = 0.0;
        int col_index = 0;
        while (std::getline(ss, item, ',')) {
            if (col_index == pz_col_index) { // Assuming pz is the third column
                pz = std::stod(item);
                min_pz = std::min(min_pz, pz);
                max_pz = std::max(max_pz, pz);
                break;
            }
            ++col_index;
        }
    }
}

class Build2DMapNode: public rclcpp::Node {
public:
    Build2DMapNode(): Node("build_2d_map_node") {
        map_pub_ = this->create_publisher<nav_msgs::msg::OccupancyGrid>("grid_map", 10);
        cloud_pub_ = this->create_publisher<sensor_msgs::msg::PointCloud2>("cloud", 10);
        this->declare_parameter<std::string>("pcl_file_path", root_dir + "src/map_manager/maps/scans.pcd");
        this->declare_parameter<std::string>("path_file_path", root_dir + "src/map_manager/maps/path.csv");
        this->declare_parameter<std::string>("map_file_path", "");
        this->declare_parameter<double>("min_pz_gain", 1.0);
        this->declare_parameter<double>("max_pz_gain", 1.0);
        this->declare_parameter<bool>("reverse", false);  // Reverse the point cloud x,z coordinates if true, since we add fast lio2 to convert x, z coordinates, this is useless now.
        this->declare_parameter<bool>("flag_build_2d_map", true);
        this->declare_parameter<double>("resolution", 0.1);
        this->declare_parameter<double>("grid_value_threshold", 50.0);
        this->declare_parameter<double>("voxel_leaf_size", 0.1);
        this->declare_parameter<double>("ror_radius", 0.05);
        this->declare_parameter<int>("ror_min_neighbors", 30);
        pcl_file_path_ = this->get_parameter("pcl_file_path").as_string();
        path_file_path_ = this->get_parameter("path_file_path").as_string();
        min_pz_gain = this->get_parameter("min_pz_gain").as_double();
        max_pz_gain = this->get_parameter("max_pz_gain").as_double();
        resolution = this->get_parameter("resolution").as_double();
        grid_value_threshold = this->get_parameter("grid_value_threshold").as_double();
        map_file_path_ = this->get_parameter("map_file_path").as_string();
        flag_build_2d_map = this->get_parameter("flag_build_2d_map").as_bool();
        reverse = this->get_parameter("reverse").as_bool();
        voxel_leaf_size = this->get_parameter("voxel_leaf_size").as_double();
        ror_radius = this->get_parameter("ror_radius").as_double();
        ror_min_neighbors = this->get_parameter("ror_min_neighbors").as_int();
        auto cvt_relative_path = [](const std::string &path) {
            if (path.empty()) return path;
            return path[0] == '/' ? path : root_dir + path;
        };
        pcl_file_path_ = cvt_relative_path(pcl_file_path_);
        path_file_path_ = cvt_relative_path(path_file_path_);
        map_file_path_ = cvt_relative_path(map_file_path_);
        if (pcl_file_path_.empty() || path_file_path_.empty()) {
            RCLCPP_ERROR(this->get_logger(), "Parameter 'pcl_file_path' and 'path_file_path'must be provided! Usage: "
                "ros2 run map_manager demo_load_pcl --ros-args -p pcl_file_path:=/path/to/your/file.pcd path_file_path:=/path/to/your/file.csv");
            rclcpp::shutdown();
            return;
        }
    }

    void build() {
        update_min_max_pz_from_path_file(path_file_path_);
        RCLCPP_INFO(this->get_logger(), "Path file loaded. Min pz: %.2f, Max pz: %.2f", min_pz, max_pz);

        RCLCPP_INFO(this->get_logger(), "Loading PCL file from: %s", pcl_file_path_.c_str());
        load_pcl_file();

        RCLCPP_INFO(this->get_logger(), "Clipping point cloud based on path file...");
        cloud_ptr_ = clip_cloud_and_reverse(cloud_ptr_, min_pz - min_pz_gain, max_pz + max_pz_gain, reverse);
        RCLCPP_INFO(this->get_logger(), "Clipping completed. Remaining points: %zu", cloud_ptr_->size());

        publish_cloud();

        if (flag_build_2d_map) {
            auto filtered_cloud = denoise_cloud();

            RCLCPP_INFO(this->get_logger(), "Building 2D map from point cloud...");
            grid_ = build_2d_map(filtered_cloud, resolution, grid_value_threshold);
            RCLCPP_INFO(this->get_logger(), "2D map built successfully with dimensions: %dx%d", grid_->info.width, grid_->info.height);

            get_save_map_file_path(grid_->info.origin.position.x, grid_->info.origin.position.y);
            save_grid_to_image();
            publish_map();
        } else if (!map_file_path_.empty()) {
            grid_ = load_grid_from_image("map", map_file_path_, resolution);
            publish_map();
        }
    }
    
    void load_pcl_file() {
        if (pcl::io::loadPCDFile<pcl::PointXYZ>(pcl_file_path_, *cloud_ptr_) == -1) {
            RCLCPP_ERROR(this->get_logger(), "Failed to load PCD file: %s", pcl_file_path_.c_str());
            rclcpp::shutdown();
            return;
        }
        RCLCPP_INFO(this->get_logger(), "Loaded PCL file with %zu points.", cloud_ptr_->size());
        RCLCPP_INFO(this->get_logger(), "width: %u, height: %u, is_dense: %s",
            cloud_ptr_->width, cloud_ptr_->height, cloud_ptr_->is_dense ? "true" : "false");
        RCLCPP_INFO(this->get_logger(), "First point: x=%.2f, y=%.2f, z=%.2f",
            cloud_ptr_->points[0].x, cloud_ptr_->points[0].y, cloud_ptr_->points[0].z);
        RCLCPP_INFO(this->get_logger(), "Sensor origin (xyz): [%.2f, %.2f, %.2f] / orientation (xyzw): [%.2f, %.2f, %.2f, %.2f]",
            cloud_ptr_->sensor_origin_.x(), cloud_ptr_->sensor_origin_.y(), cloud_ptr_->sensor_origin_.z(),
            cloud_ptr_->sensor_orientation_.x(), cloud_ptr_->sensor_orientation_.y(),
            cloud_ptr_->sensor_orientation_.z(), cloud_ptr_->sensor_orientation_.w());
    }

    void publish_cloud() {
        RCLCPP_INFO(this->get_logger(), "Publishing point cloud to /cloud ...");
        auto filtered_cloud = cloud_ptr_;
        if (voxel_leaf_size > 0) {
            filtered_cloud = apply_voxel_filter(cloud_ptr_, voxel_leaf_size);
            RCLCPP_INFO(this->get_logger(), "Point cloud downsampled with voxel leaf size: %.2f", voxel_leaf_size);
        }
        // filtered_cloud = denoise_cloud();
        sensor_msgs::msg::PointCloud2 msg;
        pcl::toROSMsg(*filtered_cloud, msg);
        msg.header.frame_id = "map";
        for (int i = 0; i < 2 && rclcpp::ok(); ++i) {
            msg.header.stamp = this->now();
            cloud_pub_->publish(msg);
            RCLCPP_INFO(this->get_logger(), "Point cloud published successfully %d times", i + 1);
            rclcpp::sleep_for(std::chrono::seconds(3));
        }
    }

    void publish_map() {
        RCLCPP_INFO(this->get_logger(), "Publishing occupancy grid map to /grid_map...");
        for (int i = 0; i < 2 && rclcpp::ok(); ++i) {
            grid_->header.stamp = this->now();
            map_pub_->publish(*grid_);
            RCLCPP_INFO(this->get_logger(), "Occupancy grid map published successfully %d times", i + 1);
            rclcpp::sleep_for(std::chrono::seconds(1));
        }
    }

    void get_save_map_file_path(float min_x, float min_y) {
        namespace fs = std::filesystem;
        fs::path p(pcl_file_path_);
        fs::path parent_path = p.parent_path();
        fs::path map_path = parent_path / ("2d_map_" +
            std::to_string(static_cast<int>(min_x * 100)) + "_" +
            std::to_string(static_cast<int>(min_y * 100)) + ".png");
        map_file_path_ = map_path.string();
    }

    void save_grid_to_image() {
        int width = grid_->info.width;
        int height = grid_->info.height;
        cv::Mat image(height, width, CV_8UC1);

        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                int index = y * width + x;
                int occupancy = grid_->data[index];
                if (occupancy == 100) {
                    image.at<uchar>(height-1-y, x) = 255; // White for occupied
                } else {
                    image.at<uchar>(height-1-y, x) = 0; // Black for unoccupied
                }
            }
        }

        cv::imwrite(map_file_path_, image);
        RCLCPP_INFO(rclcpp::get_logger("Build2DMapNode"), "Map saved to: %s", map_file_path_.c_str());
    }


private:
    std::string pcl_file_path_, path_file_path_, map_file_path_;
    point_cloud_xyz::Ptr cloud_ptr_ = std::make_shared<point_cloud_xyz>();
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr cloud_pub_;
    rclcpp::Publisher<nav_msgs::msg::OccupancyGrid>::SharedPtr map_pub_;
    std::shared_ptr<nav_msgs::msg::OccupancyGrid> grid_;

    point_cloud_xyz::Ptr denoise_cloud() {
        if (ror_radius < 0 || ror_min_neighbors < 0) {
            RCLCPP_INFO(this->get_logger(), "Skipping radius outlier removal due to invalid parameters: radius=%.2f, min_neighbors=%d", ror_radius, ror_min_neighbors);
            return cloud_ptr_;
        }
        RCLCPP_INFO(this->get_logger(), "Applying radius outlier removal with radius: %.2f, min neighbors: %d", ror_radius, ror_min_neighbors);
        auto start_time = std::chrono::steady_clock::now();
        // auto filtered_cloud = apply_voxel_filter(cloud_ptr_, 0.02, 0.02, 0.02);
        auto filtered_cloud = apply_radius_outlier_removal(cloud_ptr_, ror_radius, ror_min_neighbors);
        auto duration = std::chrono::steady_clock::now() - start_time;
        RCLCPP_INFO(this->get_logger(), "Voxel filter and radius outlier removal completed in %.2f seconds", 
            std::chrono::duration<double>(duration).count());
        RCLCPP_INFO(this->get_logger(), "Filtered point cloud size after voxel filter: %zu", filtered_cloud->size());
        return filtered_cloud;
    }
};

int main(int argc, char **argv) {
    rclcpp::init(argc, argv);
    auto node = std::make_shared<Build2DMapNode>();
    node->build();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}